#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
modbus_master_embedded.py — 从 TXT 批量驱动从站（高覆盖版，带进度）

新增：
- 实时进度打印：每 500 行或 ≥2s 打印一次 [PROGRESS]，显示 %、OK/ERR/SKIP、FPS、ETA。
- 其余逻辑不变：逐帧重连、TX/RX 打印、暖机回环等。

使用：
  1) 启动从站 (modbus_full_stack.py) 监听 0.0.0.0:1502
  2) 确认 TXT_PATH 指向你的数据集
  3) python modbus_master_embedded.py
"""

import os
import re
import time
import socket
import struct
import sys
from dataclasses import dataclass
from typing import List, Optional, Sequence, Tuple

# ===========================
# 固定配置（按需修改）
# ===========================
HOST = "127.0.0.1"
PORT = 1502
UID  = 0x01
TXT_PATH = os.path.join(os.path.dirname(__file__), "fuzz4all.txt")  # ← 改成你的数据集路径
SLEEP_BETWEEN = 0.005             # 帧间最小等待
PER_FRAME_RECONNECT = True        # True=逐帧重连；False=长连+异常时再重连
PRINT_ERR_DETAILS = True          # 打印每条失败详细原因

# 进度打印节流参数（可按需调整）
PROGRESS_EVERY_LINES = 500        # 每处理多少行打印一次进度
PROGRESS_MIN_INTERVAL = 2.0       # 两次进度打印的最小时间间隔（秒）

DEFAULT_PID = 0                   # MBAP.PID 缺省值（标准应为 0）

# ===========================
# 小工具
# ===========================
HEX2 = re.compile(r"(?:0x)?([0-9A-Fa-f]{2})")
INT_RE = re.compile(r"(-?0x[0-9A-Fa-f]+|-?\d+)")

def hexdump(b: bytes) -> str:
    return " ".join(f"{x:02X}" for x in b)

def parse_hex_any(s: str) -> Optional[bytes]:
    """从任意字符串中提取十六进制字节序列；>=7 字节认为是 MBAP+PDU。"""
    toks = HEX2.findall(s)
    if toks:
        try:
            data = bytes(int(t, 16) for t in toks)
            return data if len(data) >= 7 else None
        except ValueError:
            pass
    cleaned = "".join(ch for ch in s if ch.upper() in "0123456789ABCDEF")
    if len(cleaned) % 2 != 0:
        return None
    try:
        data = bytes(int(cleaned[i:i+2], 16) for i in range(0, len(cleaned), 2))
        return data if len(data) >= 7 else None
    except ValueError:
        return None

def parse_int(s: str, default: Optional[int]=None) -> Optional[int]:
    m = INT_RE.search(s)
    if not m:
        return default
    v = m.group(1)
    try:
        return int(v, 16) if v.lower().startswith("0x") else int(v)
    except ValueError:
        return default

def parse_list_of_uint16(s: str) -> List[int]:
    inside = s
    if "[" in s and "]" in s:
        inside = s[s.find("[")+1 : s.rfind("]")]
    parts = re.split(r"[,\s]+", inside.strip())
    out: List[int] = []
    for p in parts:
        if not p:
            continue
        try:
            v = int(p, 16) if p.lower().startswith("0x") else int(p)
            out.append(v & 0xFFFF)
        except ValueError:
            hx = HEX2.findall(p)
            if len(hx) == 2:
                out.append(((int(hx[0],16)<<8) | int(hx[1],16)) & 0xFFFF)
    return out

def _fmt_eta(sec: Optional[float]) -> str:
    if not sec or sec <= 0 or sec == float("inf"):
        return "ETA: --:--"
    m, s = divmod(int(sec), 60)
    h, m = divmod(m, 60)
    return f"ETA: {h:02d}:{m:02d}:{s:02d}"

@dataclass
class MBAP:
    tid: int
    pid: int
    length: int
    uid: int
    def pack(self) -> bytes:
        return struct.pack(">HHHB", self.tid & 0xFFFF, self.pid & 0xFFFF, self.length & 0xFFFF, self.uid & 0xFF)
    @staticmethod
    def unpack(b: bytes) -> "MBAP":
        tid, pid, length, uid = struct.unpack(">HHHB", b)
        return MBAP(tid, pid, length, uid)

# ===========================
# 主站实现
# ===========================
class ModbusMaster:
    def __init__(self, host: str, port: int, uid: int) -> None:
        self.host = host
        self.port = port
        self.uid  = uid
        self.sock: Optional[socket.socket] = None
        self.tid  = 0

    def connect(self) -> None:
        self.sock = socket.create_connection((self.host, self.port), timeout=3.0)
        self.sock.settimeout(2.0)  # 读写超时，避免挂死

    def close(self) -> None:
        if self.sock is not None:
            try:
                self.sock.close()
            finally:
                self.sock = None

    def _ensure(self) -> None:
        if self.sock is None:
            self.connect()

    def _read_exact(self, n: int) -> bytes:
        buf = bytearray()
        while len(buf) < n:
            chunk = self.sock.recv(n - len(buf))  # type: ignore[arg-type]
            if not chunk:
                raise ConnectionError("server closed")
            buf.extend(chunk)
        return bytes(buf)

    # ---- PDU（自动补MBAP）----
    def txrx_pdu(self, pdu: bytes, pid: int = DEFAULT_PID) -> bytes:
        self._ensure()
        self.tid = (self.tid + 1) & 0xFFFF
        mbap = MBAP(self.tid, pid, 1 + len(pdu), UID).pack()
        raw = mbap + pdu
        print("[TX] PDU:", hexdump(raw))
        self.sock.sendall(raw)  # type: ignore[arg-type]
        hdr = self._read_exact(7)
        mbap_r = MBAP.unpack(hdr)
        body = self._read_exact(max(0, mbap_r.length - 1))
        print("[RX] PDU:", hexdump(hdr + body))
        return body

    # ---- 原始整帧（MBAP+PDU）----
    def send_frame_raw(self, raw_frame: bytes) -> bytes:
        self._ensure()
        print("[TX] RAW:", hexdump(raw_frame))
        self.sock.sendall(raw_frame)  # type: ignore[arg-type]
        hdr = self._read_exact(7)
        mbap_r = MBAP.unpack(hdr)
        body = self._read_exact(max(0, mbap_r.length - 1))
        print("[RX] RAW:", hexdump(hdr + body))
        return body

    # ---- 便捷操作 ----
    def write_multiple_regs(self, start: int, regs: Sequence[int]) -> None:
        qty = len(regs); bc = 2 * qty
        pdu = struct.pack(">BHHB", 0x10, start, qty, bc) + struct.pack(">" + "H"*qty, *[r & 0xFFFF for r in regs])
        rsp = self.txrx_pdu(pdu)
        if rsp and rsp[0] == 0x90:
            raise RuntimeError(f"exception {rsp[1]}")
        if len(rsp) != 5 or rsp[0] != 0x10:
            raise RuntimeError("invalid response for 0x10")

    def read_holding(self, start: int, qty: int) -> List[int]:
        pdu = struct.pack(">BHH", 0x03, start, qty)
        rsp = self.txrx_pdu(pdu)
        if rsp and rsp[0] == 0x83:
            raise RuntimeError(f"exception {rsp[1]}")
        if rsp[0] != 0x03:
            raise RuntimeError(f"unexpected fc in response: {rsp[0]:#x}")
        bc = rsp[1]
        if len(rsp) != 2 + bc:
            raise RuntimeError("byte count mismatch")
        return list(struct.unpack(">" + "H" * (bc//2), rsp[2:2+bc]))

# ===========================
# 文本行解析（A/B/C 三路兜底）
# ===========================
def build_mbap_pdu_frame_from_pdu_hex(pdu_hex_line: str, pid: Optional[int]=None, uid: Optional[int]=None, tid: Optional[int]=None) -> Optional[bytes]:
    pdu = parse_hex_any(pdu_hex_line)
    if not pdu:
        return None
    _pid = DEFAULT_PID if pid is None else pid
    _uid = UID if uid is None else uid
    _tid = 1 if tid is None else tid
    mbap = MBAP(_tid & 0xFFFF, _pid & 0xFFFF, 1 + len(pdu), _uid & 0xFF).pack()
    return mbap + pdu

def parse_line_to_frame(line: str) -> Tuple[Optional[bytes], str]:
    """
    返回 (frame_bytes, reason)，当 frame_bytes 为 None 时 reason 表示跳过原因。
    优先级：
      A) 原始整帧（MBAP+PDU）十六进制（可无空格）
      B) "pdu= ..." 十六进制 → 自动包MBAP
      C) 键值对（fc/start/qty/vals/pid/uid/tid）→ 组PDU并包MBAP（支持fc=3/16）
    """
    s = line.strip()
    if not s or s.startswith(("#","//",";")):
        return None, "comment/empty"

    # A) 原始整帧
    hx = parse_hex_any(s)
    if hx is not None:
        return hx, "A:raw_hex"

    # B) pdu= ...
    low = s.lower()
    if "pdu" in low:
        m = re.search(r"pdu\s*=\s*(.*)$", s, flags=re.IGNORECASE)
        if m:
            pid = parse_int(re.findall(r"pid\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"pid\s*=", s, flags=re.IGNORECASE) else None
            uid = parse_int(re.findall(r"uid\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"uid\s*=", s, flags=re.IGNORECASE) else None
            tid = parse_int(re.findall(r"tid\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"tid\s*=", s, flags=re.IGNORECASE) else None
            frame = build_mbap_pdu_frame_from_pdu_hex(m.group(1), pid=pid, uid=uid, tid=tid)
            return (frame, "B:pdu_plus_mbap") if frame else (None, "bad pdu hex")

    # C) 键值对
    if "fc" in low:
        fc  = parse_int(re.findall(r"fc\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"fc\s*=", s, flags=re.IGNORECASE) else None
        pid = parse_int(re.findall(r"pid\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"pid\s*=", s, flags=re.IGNORECASE) else DEFAULT_PID
        uid = parse_int(re.findall(r"uid\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"uid\s*=", s, flags=re.IGNORECASE) else UID
        tid = parse_int(re.findall(r"tid\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"tid\s*=", s, flags=re.IGNORECASE) else 1

        if fc is None:
            return None, "kv: missing fc"

        if fc == 3:
            start = parse_int(re.findall(r"start\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"start\s*=", s, flags=re.IGNORECASE) else 0
            qty   = parse_int(re.findall(r"qty\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"qty\s*=", s, flags=re.IGNORECASE) else 1
            pdu = struct.pack(">BHH", 0x03, start or 0, qty or 1)
            mbap = MBAP(tid & 0xFFFF, pid & 0xFFFF, 1 + len(pdu), uid & 0xFF).pack()
            return mbap + pdu, "C:kv_fc03"

        if fc in (16, 0x10):
            start = parse_int(re.findall(r"start\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"start\s*=", s, flags=re.IGNORECASE) else 0
            vals_m = re.search(r"(vals|values|data|regs)\s*=\s*(\[.*\]|.+)$", s, flags=re.IGNORECASE)
            regs: List[int] = parse_list_of_uint16(vals_m.group(2)) if vals_m else []
            qty = parse_int(re.findall(r"qty\s*=\s*([^\s,;]+)", s, flags=re.IGNORECASE)[0]) if re.search(r"qty\s*=", s, flags=re.IGNORECASE) else (len(regs) if regs else 0)
            if qty is None or qty <= 0:
                if not regs:
                    regs = [0x0000]
                qty = len(regs)
            if len(regs) < qty:
                regs = regs + [0x0000] * (qty - len(regs))
            bc = 2 * qty
            pdu = struct.pack(">BHHB", 0x10, start or 0, qty, bc) + struct.pack(">" + "H"*qty, *regs[:qty])
            mbap = MBAP(tid & 0xFFFF, pid & 0xFFFF, 1 + len(pdu), uid & 0xFF).pack()
            return mbap + pdu, "C:kv_fc10"

        return None, f"kv: unsupported fc={fc}"

    return None, "unrecognized"

# ===========================
# 逐行读取并发送（带实时进度）
# ===========================
class BatchStats:
    def __init__(self) -> None:
        self.total_lines = 0
        self.sent_ok = 0
        self.sent_err = 0
        self.skipped = 0
        self.skip_reasons = {}

    def add_skip(self, reason: str) -> None:
        self.skipped += 1
        self.skip_reasons[reason] = self.skip_reasons.get(reason, 0) + 1

def send_frames_from_txt(master: ModbusMaster, txt_path: str) -> None:
    if not os.path.exists(txt_path):
        raise FileNotFoundError(f"TXT not found: {txt_path}")

    # 预统计总行数（用于百分比与 ETA）
    with open(txt_path, "r", encoding="utf-8", errors="ignore") as fcnt:
        total_lines = sum(1 for _ in fcnt)

    stat = BatchStats()
    start_t = time.monotonic()
    last_progress_t = start_t

    print(f"\n=== 读取并发送 TXT：{os.path.basename(txt_path)} ===")
    print(f"[INFO] total lines (including comments/blanks): {total_lines}")

    def maybe_print_progress(li: int):
        nonlocal last_progress_t
        now = time.monotonic()
        if (li % PROGRESS_EVERY_LINES == 0) or (now - last_progress_t >= PROGRESS_MIN_INTERVAL):
            elapsed = max(1e-6, now - start_t)
            fps = stat.sent_ok / elapsed
            pct = (li / total_lines * 100.0) if total_lines > 0 else 0.0
            remaining = max(0, total_lines - li)
            eta = (remaining / fps) if fps > 0 else None
            print(f"[PROGRESS] {li}/{total_lines} ({pct:6.2f}%)  OK={stat.sent_ok}  ERR={stat.sent_err}  "
                  f"SKIP={stat.skipped}  FPS={fps:6.1f}  {_fmt_eta(eta)}")
            last_progress_t = now

    with open(txt_path, "r", encoding="utf-8", errors="ignore") as f:
        for li, line in enumerate(f, 1):
            stat.total_lines += 1
            frame, info = parse_line_to_frame(line)
            if frame is None:
                stat.add_skip(info)
                maybe_print_progress(li)
                continue

            try:
                if PER_FRAME_RECONNECT:
                    master.close()
                    master.connect()
                master.send_frame_raw(frame)
                stat.sent_ok += 1
            except Exception as e:
                stat.sent_err += 1
                if PRINT_ERR_DETAILS:
                    print(f"[ERROR] line {li}: {e!r}")
                try:
                    master.close()
                    time.sleep(0.02)
                    master.connect()
                except Exception:
                    pass
            finally:
                maybe_print_progress(li)
                time.sleep(SLEEP_BETWEEN)

    # 结束时给一次 100% 的快照
    elapsed = max(1e-6, time.monotonic() - start_t)
    fps = stat.sent_ok / elapsed
    print(f"[PROGRESS] {stat.total_lines}/{total_lines} (100.00%)  OK={stat.sent_ok}  ERR={stat.sent_err}  "
          f"SKIP={stat.skipped}  FPS={fps:6.1f}  {_fmt_eta(0)}")

    print("\n=== 汇总 ===")
    print(f"lines={stat.total_lines} ok={stat.sent_ok} err={stat.sent_err} skipped={stat.skipped}")
    print(f"skip_breakdown={stat.skip_reasons}")

# ===========================
# 主流程
# ===========================
def main() -> None:
    master = ModbusMaster(HOST, PORT, UID)

    # 1) 暖机：0x10 → 0x03
    try:
        master.connect()
        print("=== 正常写/读回环（0x10 -> 0x03） ===")
        master.write_multiple_regs(0, [0x1111, 0x2222, 0x3333])
        time.sleep(SLEEP_BETWEEN)
        try:
            master.read_holding(0, 6)
        except Exception as e:
            print(f"[WARN] warm read exception: {e!r}")
        time.sleep(SLEEP_BETWEEN)
    finally:
        master.close()

    # 2) 批量发送
    if not PER_FRAME_RECONNECT:
        try:
            master.connect()
        except Exception as e:
            print(f"[WARN] initial connect failed: {e!r}")

    try:
        send_frames_from_txt(master, TXT_PATH)
    finally:
        master.close()

    print("\n=== 完成 ===")

if __name__ == "__main__":
    main()
